#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
角速度预测工具

该程序用于利用已训练的角速度模型进行预测，可以预测指定命令值对应的角速度，
或查找产生目标角速度的命令值。支持批量预测、生成预测表格和绘制预测图表。
"""

import os  # 用于文件路径处理
import sys  # 用于系统级操作
import numpy as np  # 用于数值计算
import pandas as pd  # 用于数据处理
import matplotlib as mpl  # 绘图库
# 设置Agg后端，确保在无GUI环境下也能正常生成图像
mpl.use('Agg')
import matplotlib.pyplot as plt  # 用于绘制图表
from datetime import datetime  # 用于生成时间戳
import joblib  # 用于加载模型
import argparse  # 用于解析命令行参数
from sklearn.linear_model import LinearRegression  # 用于创建默认模型
from sklearn.preprocessing import PolynomialFeatures  # 用于创建多项式特征转换器

# 数据路径设置：指定模型文件路径
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_FILES = {
    'rtk': os.path.join(BASE_DIR, 'RTK角速度实验_angular_velocity_model.joblib'),
    'imu': os.path.join(BASE_DIR, 'IMU角速度实验_angular_velocity_model.joblib'),
    'combined': os.path.join(BASE_DIR, 'combined_angular_velocity_model.joblib')
}
POLY_FILES = {
    'rtk': os.path.join(BASE_DIR, 'RTK角速度实验_angular_velocity_poly_features.joblib'),
    'imu': os.path.join(BASE_DIR, 'IMU角速度实验_angular_velocity_poly_features.joblib'),
    'combined': os.path.join(BASE_DIR, 'combined_angular_velocity_poly_features.joblib')
}

# 输出配置：定义图表尺寸、分辨率和输出路径等
OUTPUT_CONFIG = {
    'figsize': (10, 6),  # 图表大小
    'dpi': 100,  # 图表分辨率
    'save_path': 'predictions',  # 保存目录
    'prediction_file': 'angular_velocity_predictions.csv'  # 预测结果文件名
}

# 创建默认模型函数
def create_default_model(model_type):
    """创建一个默认的角速度预测模型
    
    当找不到预训练模型文件时，创建一个简单的线性模型作为替代。
    这个模型将命令值映射到角速度的简单关系，仅用于演示目的。
    
    参数:
        model_type: 模型类型名称
        
    返回:
        model: 简单线性回归模型
        poly: 多项式特征转换器
    """
    print(f"警告: 找不到{model_type}模型文件，创建默认示例模型")
    
    # 创建一个简单的线性回归模型
    model = LinearRegression()
    
    # 创建3次多项式特征转换器
    poly = PolynomialFeatures(degree=3)
    
    # 生成一些示例数据
    commands = np.array([500, 1000, 1500, 2000]).reshape(-1, 1)
    
    # 根据模型类型设置不同的模拟角速度关系
    if model_type == 'rtk':
        # RTK模型示例：线性关系
        angular_velocities = np.array([-30, -10, 10, 30])
    elif model_type == 'imu':
        # IMU模型示例：略带非线性
        angular_velocities = np.array([-35, -12, 12, 35])
    else:  # combined
        # 组合模型示例：更精确的关系
        angular_velocities = np.array([-32, -11, 11, 32])
    
    # 准备多项式特征
    X_poly = poly.fit_transform(commands)
    
    # 训练模型
    model.fit(X_poly, angular_velocities)
    
    # 保存模型和多项式特征转换器
    os.makedirs(os.path.dirname(MODEL_FILES[model_type]), exist_ok=True)
    joblib.dump(model, MODEL_FILES[model_type])
    joblib.dump(poly, POLY_FILES[model_type])
    
    print(f"已创建并保存默认{model_type}模型")
    
    return model, poly

class AngularVelocityPredictor:
    """角速度预测器类
    
    封装了角速度预测所需的所有功能，包括模型加载、单值预测、批量预测、
    生成预测表格、绘制预测图表，以及查找特定角速度对应的命令值。
    """
    
    def __init__(self, model_type='combined'):
        """初始化预测器
        
        创建预测器实例，加载指定类型的预训练模型和多项式特征转换器。
        
        参数:
            model_type: 模型类型，可选 'rtk', 'imu', 'combined'，
                       分别对应RTK数据训练的模型、IMU数据训练的模型和合并数据训练的模型
        """
        self.model_type = model_type
        self.model = None  # 预测模型
        self.poly = None  # 多项式特征转换器
        self.load_model()  # 初始化时自动加载模型
        
    def load_model(self):
        """加载模型文件
        
        根据初始化时指定的model_type，加载对应的预训练模型和多项式特征转换器。
        如果找不到模型文件或加载出错，会抛出异常。
        
        异常:
            FileNotFoundError: 找不到模型文件时抛出
            RuntimeError: 加载模型出错时抛出
        """
        model_file = MODEL_FILES.get(self.model_type)
        poly_file = POLY_FILES.get(self.model_type)
        
        # 检查模型文件是否存在
        model_exists = os.path.exists(model_file) if model_file else False
        poly_exists = os.path.exists(poly_file) if poly_file else False
        
        # 如果模型文件不存在，则创建默认模型
        if not model_exists or not poly_exists:
            self.model, self.poly = create_default_model(self.model_type)
            return
        
        try:
            # 加载模型和多项式特征转换器
            self.model = joblib.load(model_file)
            self.poly = joblib.load(poly_file)
            print(f"成功加载 {self.model_type} 模型")
        except Exception as e:
            print(f"加载模型出错: {e}，将创建默认模型")
            self.model, self.poly = create_default_model(self.model_type)
    
    def predict(self, command):
        """预测单个命令值对应的角速度
        
        将单个命令值转换为多项式特征，然后使用模型预测对应的角速度。
        
        参数:
            command: 电机命令值，通常在500-2000之间
            
        返回:
            预测的角速度 (度/秒)
            
        异常:
            RuntimeError: 模型未加载时抛出
        """
        if self.model is None or self.poly is None:
            raise RuntimeError("模型未加载，请先加载模型")
        
        # 将命令值转换为特征矩阵格式
        X = np.array([command]).reshape(-1, 1)
        # 应用多项式变换
        X_poly = self.poly.transform(X)
        
        # 预测角速度并返回单个值
        angular_velocity = self.model.predict(X_poly)[0]
        return angular_velocity
    
    def predict_batch(self, commands):
        """批量预测命令值对应的角速度
        
        将多个命令值转换为多项式特征，然后使用模型批量预测对应的角速度。
        
        参数:
            commands: 电机命令值列表或数组
            
        返回:
            预测的角速度列表 (度/秒)
            
        异常:
            RuntimeError: 模型未加载时抛出
        """
        if self.model is None or self.poly is None:
            raise RuntimeError("模型未加载，请先加载模型")
        
        # 将命令值列表转换为特征矩阵格式
        X = np.array(commands).reshape(-1, 1)
        # 应用多项式变换
        X_poly = self.poly.transform(X)
        
        # 批量预测角速度
        angular_velocities = self.model.predict(X_poly)
        return angular_velocities
    
    def generate_prediction_table(self, min_val=500, max_val=2000, step=100):
        """生成预测表格
        
        在指定范围内以给定步长生成命令值，然后预测对应的角速度，返回预测结果表格。
        
        参数:
            min_val: 最小命令值，默认为500
            max_val: 最大命令值，默认为2000
            step: 命令值步长，默认为100
            
        返回:
            预测表格 (DataFrame)，包含命令值和对应的预测角速度
        """
        # 生成等间隔的命令值序列
        commands = list(range(min_val, max_val + 1, step))
        # 批量预测角速度
        angular_velocities = self.predict_batch(commands)
        
        # 创建数据表，将命令值和预测结果组合
        df = pd.DataFrame({
            'Command': commands,
            'AngularVelocity': angular_velocities
        })
        
        return df
    
    def save_prediction_table(self, df, output_dir=None):
        """保存预测表格
        
        将预测结果表格保存为CSV文件，文件名包含模型类型和时间戳。
        
        参数:
            df: 预测表格 (DataFrame)
            output_dir: 输出目录，如果为None则使用默认目录
            
        返回:
            保存的文件完整路径
        """
        # 确定输出目录
        if output_dir is None:
            output_dir = os.path.join(BASE_DIR, OUTPUT_CONFIG['save_path'])
        
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        # 获取当前时间作为文件名的一部分
        current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 拼接完整的输出文件路径
        output_file = os.path.join(output_dir, f"{self.model_type}_predictions_{current_time}.csv")
        # 保存CSV文件
        df.to_csv(output_file, index=False)
        
        return output_file
    
    def plot_predictions(self, df, output_dir=None, show_equation=True):
        """绘制预测图表
        
        根据预测表格绘制命令值与角速度的关系图，可选显示拟合方程，并保存为图像文件。
        
        参数:
            df: 预测表格 (DataFrame)
            output_dir: 输出目录，如果为None则使用默认目录
            show_equation: 是否在图表上显示拟合方程
            
        返回:
            保存的图表文件完整路径
        """
        # 创建指定大小和分辨率的图表
        plt.figure(figsize=OUTPUT_CONFIG['figsize'], dpi=OUTPUT_CONFIG['dpi'])
        
        # 绘制预测点：散点图
        plt.scatter(df['Command'], df['AngularVelocity'], color='blue', label='Predicted Points')
        
        # 绘制拟合曲线：生成更密集的点以获得平滑曲线
        x = np.linspace(df['Command'].min(), df['Command'].max(), 100)
        x_poly = self.poly.transform(x.reshape(-1, 1))
        y_pred = self.model.predict(x_poly)
        
        plt.plot(x, y_pred, 'r-', label='Fitted Curve')
        
        # 显示拟合方程：将模型系数转换为数学表达式
        if show_equation and hasattr(self.model, 'coef_') and hasattr(self.model, 'intercept_'):
            degree = self.poly.degree
            coefs = self.model.coef_[0]
            intercept = self.model.intercept_
            
            # 构建方程字符串
            equation = f"y = {intercept:.4f}"
            for i in range(1, len(coefs)):
                if coefs[i] >= 0:
                    equation += f" + {coefs[i]:.6f}x^{i}"
                else:
                    equation += f" - {abs(coefs[i]):.6f}x^{i}"
            
            # 在图表左上角添加方程文本
            plt.text(0.05, 0.95, equation, transform=plt.gca().transAxes, 
                    fontsize=10, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        
        # 添加标签和标题
        plt.xlabel('Command Value')
        plt.ylabel('Angular Velocity (deg/s)')
        plt.title(f'{self.model_type.upper()} Angular Velocity Prediction')
        plt.grid(True)
        plt.legend()
        
        # 确定输出目录
        if output_dir is None:
            output_dir = os.path.join(BASE_DIR, OUTPUT_CONFIG['save_path'])
        
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        # 获取当前时间作为文件名的一部分
        current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 保存图表为PNG文件
        output_file = os.path.join(output_dir, f"{self.model_type}_prediction_plot_{current_time}.png")
        plt.savefig(output_file)
        plt.close()  # 关闭图表释放资源
        
        return output_file
    
    def find_command_for_angular_velocity(self, target_velocity, min_val=500, max_val=2000, num_points=1000):
        """查找给定角速度对应的命令值
        
        通过密集采样在给定范围内生成多个命令值，预测它们的角速度，
        然后找出与目标角速度最接近的命令值。
        
        参数:
            target_velocity: 目标角速度 (度/秒)
            min_val: 搜索范围最小命令值
            max_val: 搜索范围最大命令值
            num_points: 搜索点数，越大精度越高但速度越慢
            
        返回:
            最接近目标角速度的命令值
        """
        # 生成密集的命令值数组用于搜索
        commands = np.linspace(min_val, max_val, num_points)
        
        # 预测这些命令值对应的角速度
        angular_velocities = self.predict_batch(commands)
        
        # 计算每个角速度与目标值的绝对差距
        differences = np.abs(angular_velocities - target_velocity)
        
        # 找到差距最小的索引
        closest_idx = np.argmin(differences)
        
        # 返回对应的命令值
        return commands[closest_idx]

def parse_arguments():
    """解析命令行参数
    
    设置并解析命令行参数，支持选择模型类型、预测模式、命令值范围等。
    
    返回:
        解析后的命令行参数对象
    """
    parser = argparse.ArgumentParser(description='角速度预测工具')
    
    parser.add_argument('--model', type=str, default='combined', choices=['rtk', 'imu', 'combined'],
                        help='使用的模型类型 (默认: combined)')
    
    parser.add_argument('--predict', type=float, nargs='+',
                        help='预测特定命令值的角速度')
    
    parser.add_argument('--find-command', type=float,
                        help='查找给定角速度对应的命令值')
    
    parser.add_argument('--min-command', type=int, default=500,
                        help='最小命令值 (默认: 500)')
    
    parser.add_argument('--max-command', type=int, default=2000,
                        help='最大命令值 (默认: 2000)')
    
    parser.add_argument('--step', type=int, default=100,
                        help='命令值步长 (默认: 100)')
    
    parser.add_argument('--generate-table', action='store_true',
                        help='生成预测表格')
    
    parser.add_argument('--plot', action='store_true',
                        help='绘制预测图表')
    
    parser.add_argument('--output-dir', type=str, default='',
                        help='输出目录 (默认为脚本所在目录下的predictions文件夹)')
    
    return parser.parse_args()

def print_banner():
    """打印命令行界面横幅
    
    在程序启动时打印美观的横幅，提升用户体验并显示程序名称。
    """
    banner = """
=================================================================
                   角速度预测工具 (Angular Velocity Prediction Tool)
=================================================================
    """
    print(banner)

def main():
    """主函数
    
    程序入口点，执行整个预测流程，包括参数解析、模型加载、
    预测执行以及结果展示/保存等。
    
    返回:
        状态码，0表示正常执行，1表示执行出错
    """
    # 打印程序横幅
    print_banner()
    
    # 解析命令行参数
    args = parse_arguments()
    
    try:
        # 创建预测器实例并加载模型
        predictor = AngularVelocityPredictor(model_type=args.model)
        
        # 设置输出目录
        output_dir = args.output_dir if args.output_dir else os.path.join(BASE_DIR, OUTPUT_CONFIG['save_path'])
        
        # 如果指定了预测命令值选项
        if args.predict:
            print("\n命令值预测结果:")
            print("-" * 30)
            print("命令值\t角速度(度/秒)")
            print("-" * 30)
            
            # 遍历并预测每个命令值
            for command in args.predict:
                angular_velocity = predictor.predict(command)
                print(f"{command}\t{angular_velocity:.4f}")
                
            print("-" * 30)
        
        # 如果指定了查找命令值选项
        if args.find_command is not None:
            target_velocity = args.find_command
            # 查找能产生目标角速度的命令值
            command = predictor.find_command_for_angular_velocity(
                target_velocity, min_val=args.min_command, max_val=args.max_command
            )
            print(f"\n目标角速度 {target_velocity:.4f} 度/秒对应的命令值约为: {command:.2f}")
        
        # 如果需要生成预测表格
        if args.generate_table:
            # 生成预测表格
            df = predictor.generate_prediction_table(
                min_val=args.min_command, max_val=args.max_command, step=args.step
            )
            
            # 保存预测表格
            output_file = predictor.save_prediction_table(df, output_dir=output_dir)
            print(f"\n预测表格已保存到: {output_file}")
            
            # 如果需要绘制图表
            if args.plot:
                plot_file = predictor.plot_predictions(df, output_dir=output_dir)
                print(f"预测图表已保存到: {plot_file}")
        
        print("\n预测完成!")
        
    except Exception as e:
        # 捕获并显示任何异常
        print(f"\n预测出错: {e}")
        return 1  # 返回错误状态码
    
    return 0  # 返回成功状态码

if __name__ == "__main__":
    # 调用主函数并将返回值传递给sys.exit
    sys.exit(main()) 